package promote.brozen.algorithm.tree.balance.bpt;

import promote.brozen.algorithm.tree.Tree;

import java.util.function.Consumer;

/**
 * @author Brozen
 * @date 2020/1/17 5:26 PM
 */
public class BPlusTree<K extends Comparable<K>> implements Tree<K> {

    private BPlusTreeNode<K> root;

    private final int degree;

    public BPlusTree() {
        this(3);
    }

    public BPlusTree(int degree) {
        this.degree = degree;
    }

    @Override
    public void add(K key) {
        // 如果root为空，设置root
        if (this.root == null) {
            this.root = new BPlusTreeNode<>(key);
            return;
        }

        // 1. 添加到树
        BPlusTreeDataChainNode<K> addedNode = doAdd(this.root, key);

        // 2. 如果添加后，树节点数超过degree-1，需要进行平衡
        BPlusTreeNode<K> treeNode = addedNode.getNode();
        if (treeNode.getKeySize() >= this.degree) {
            doAddBalance(treeNode);
        }
    }

    private BPlusTreeDataChainNode<K> doAdd(BPlusTreeNode<K> node, K key) {
        // 1. 在当前树节点链表中，找到插入节点的前一个key节点位置
        BPlusTreeDataChainNode<K> keyNode = node.getHead().getNext();
        int result = 0;
        while (keyNode != null) {
            result = keyNode.getKey().compareTo(key);
            if (result > 0) {
                // 找到了要插入节点的后一个节点位置
                break;
            } else {
                keyNode = keyNode.getNext().getNext();
            }
        }

        if (!node.isLeaf()) {
            // 2. 如果不是叶子节点，需要继续向下一层查询
            BPlusTreeChildChainNode<K> childPointer = keyNode == null ? node.getTail() : keyNode.getPrev();
            return doAdd(childPointer.getChild(), key);
        } else {
            // 3. 如果是叶子节点，需要在此节点链表中插入
            BPlusTreeChildChainNode<K> head = node.getHead();
            BPlusTreeChildChainNode<K> tail = node.getTail();
            BPlusTreeDataChainNode<K> newDataNode = new BPlusTreeDataChainNode<>(node, key);
            BPlusTreeChildChainNode<K> newChildNode = new BPlusTreeChildChainNode<>(node);
            if (keyNode == null) {
                // 3.1 如果keyNode为null，说明需要插入到尾节点
                tail.setNext(newDataNode);
                newDataNode.setPrev(tail);
                newDataNode.setNext(newChildNode);
                newChildNode.setPrev(newDataNode);
                node.setTail(newChildNode);
                node.increaseKeySize();
            } else {
                BPlusTreeDataChainNode<K> prevKey = keyNode.getPrev().getPrev();
                if(prevKey != null) {
                    // 3.3 如果keyNode前一个key存在，说明不是插入到头节点
                    if (prevKey.getKey().compareTo(key) == 0) {
                        // 3.3.1 如果keyNode的前一个key与插入key值相等，则需要更新
                        prevKey.setKey(key);
                    } else {
                        // 3.3.2 正常插入到keyNode之前
                        newDataNode.setPrev(keyNode.getPrev());
                        keyNode.getPrev().setNext(newDataNode);
                        newDataNode.setNext(newChildNode);
                        newChildNode.setPrev(newDataNode);
                        newChildNode.setNext(keyNode);
                        keyNode.setPrev(newChildNode);
                        node.increaseKeySize();
                    }
                } else {
                    // 3.2 如果keyNode的前一个key为null，说明需要插入到头结点
                    newChildNode.setNext(newDataNode);
                    newDataNode.setPrev(newChildNode);
                    newDataNode.setNext(head);
                    head.setPrev(newDataNode);
                    node.setHead(newChildNode);
                    node.increaseKeySize();
                }
            }

            return newDataNode;
        }
    }

    /**
     * 新增时平衡操作，进入的node，其keySize都是 >= tree.degree的
     * 递归临界条件：分裂后，平衡节点的父节点keySize < tree.degree
     */
    private void doAddBalance(BPlusTreeNode<K> node) {

        // 1. 找到要分裂后上升的key节点
        int keySize = node.getKeySize();
        int upIndex = keySize / 2;
        BPlusTreeDataChainNode<K> upKeyNode = node.getDataChainNode(upIndex);

        BPlusTreeChildChainNode<K> parentPointer = node.getParent();
        BPlusTreeDataChainNode<K> parentPointerNext = parentPointer == null ? null : parentPointer.getNext();
        BPlusTreeNode<K> parentTreeNode = parentPointer == null ? null : parentPointer.getNode();

        // 2. 将节点分裂成两部分，左边的在原来节点，右边的用新的节点，并连接两部分的链表
        // 2.1 左边部分在原来节点，重设左边部分的
        BPlusTreeChildChainNode<K> leftNewTail = upKeyNode.getPrev();
        BPlusTreeChildChainNode<K> leftOldTail = node.getTail();
        leftNewTail.setNext(null);
        node.setTail(leftNewTail);
        node.setKeySize(upIndex);
        // 2.2 右边部分用新的节点存储，并且isLeaf与平衡节点相同
        BPlusTreeNode<K> rightNode = new BPlusTreeNode<>();
        rightNode.setLeaf(node.isLeaf());
        BPlusTreeChildChainNode<K> rightHead;
        if (node.isLeaf()) {
            // 如果是叶子节点，upKeyNode是保持在叶子节点上，实际上升的是一个新的节点，但拥有与upKeyNode一样的key
            rightHead = new BPlusTreeChildChainNode<>(rightNode);
            rightHead.setNext(upKeyNode);
            upKeyNode.setPrev(rightHead);
            rightNode.setKeySize(keySize - upIndex);
            node.setNext(rightNode);
            upKeyNode = new BPlusTreeDataChainNode<>(rightNode, upKeyNode.getKey());
        } else {
            // 如果是非叶子节点，upKeyNode是要真正的上升的
            rightHead = upKeyNode.getNext();
            rightNode.setKeySize(keySize - upIndex - 1);
        }
        rightHead.setPrev(null);
        rightNode.setHead(rightHead);
        rightNode.setTail(leftOldTail);
        BPlusTreeChildChainNode<K> cn = rightHead;
        while (cn != null) {
            cn.setNode(rightNode);
            if (cn.getNext() != null) {
                cn.getNext().setNode(rightNode);
                cn = cn.getNext().getNext();
            } else {
                break;
            }
        }

        // 3. 将upKeyNode放入父节点，前面的child指针指向平衡节点，后面的child指针指向新建的节点
        if (parentTreeNode == null) {
            // 3.1 如果父节点不存在，说明是根节点分裂，需要创建一个新的根节点
            this.root = parentTreeNode = new BPlusTreeNode<>(upKeyNode.getKey());
            BPlusTreeChildChainNode<K> parentTreeNodeHead = parentTreeNode.getHead();
            BPlusTreeChildChainNode<K> parentTreeNodeTail = parentTreeNode.getTail();
            parentTreeNodeHead.setChild(node);
            parentTreeNodeTail.setChild(rightNode);
            parentTreeNode.setLeaf(false);
            node.setParent(parentTreeNodeHead);
            rightNode.setParent(parentTreeNodeTail);
        } else {
            // 3.2 父节点存在，非根节点裂开
            BPlusTreeChildChainNode<K> rightNodeParent = new BPlusTreeChildChainNode<>(parentTreeNode);
            rightNodeParent.setChild(rightNode);
            rightNode.setParent(rightNodeParent);
            parentPointer.setNext(upKeyNode);
            upKeyNode.setPrev(parentPointer);
            upKeyNode.setNode(parentTreeNode);
            upKeyNode.setNext(rightNodeParent);
            rightNodeParent.setPrev(upKeyNode);
            rightNodeParent.setNext(parentPointerNext);
            if (parentPointerNext != null) {
                parentPointerNext.setPrev(rightNodeParent);
            } else {
                // 上升节点之后没有其他key，说明上升节点是插入到了父节点的尾巴上，需要重设tail
                parentTreeNode.setTail(rightNodeParent);
            }
            parentTreeNode.increaseKeySize();
        }

        // 4. 如果父节点keySize >= tree.degree，则将父节点带入平衡操作递归处理
        if (parentTreeNode.getKeySize() >= this.degree) {
            doAddBalance(parentTreeNode);
        }
    }

    @Override
    public void remove(K key) {

    }

    @Override
    public K search(K key) {
        return null;
    }

    public void walk(Consumer<K> consumer) {
        BPlusTreeChildChainNode<K> node = this.root.getHead();
        while (node.getChild() != null) {
            node = node.getChild().getHead();
        }

        while (node != null) {
            BPlusTreeDataChainNode<K> dataNode = node.getNext();
            if (dataNode == null) {
                BPlusTreeNode<K> nextTreeNode = node.getNode().getNext();
                node = nextTreeNode == null ? null : nextTreeNode.getHead();
            } else {
                consumer.accept(dataNode.getKey());
                node = dataNode.getNext();
            }
        }

    }
}
